Shader "Hidden/VPR/TAA"
{
    HLSLINCLUDE
    
    #include "Packages/com.unity.render-pipelines.universal/ShaderLibrary/Core.hlsl"
    TEXTURE2D(_SourceTex);
    
    TEXTURE2D(_CameraDepthTex);
    float4 _CameraDepthTex_TexelSize;
    SamplerState sampler_CameraDepthTex;
    SamplerState sampler_PointClamp;
    SamplerState sampler_LinearClamp;
    
    TEXTURE2D(_HistoryBuffer);
    SamplerState sampler_HistoryBuffer;
    float4 _SourceTex_TexelSize;  
    int _IgnoreHistory;
    float _BlendAlpha;
    struct appdata
    {
        float4 vertex : POSITION;
        float2 uv : TEXCOORD0;
    };

    struct v2f
    {
        
        float4 vertex : SV_POSITION;
        float2 uv : TEXCOORD0; //存储两个uv坐标
        
    };

    float4x4 _PrevViewProjM_TAA;
    float4x4 _I_P_Current_jittered;
    
    float4x4 _I_V_Current_jittered;
    
    
    float2 PreviousOffset;
    float2 CurrentOffset;
    //float2 ScreenSize;

    
    float Luminance(float3 color)
    {
        return 0.25 * color.r + 0.5 * color.g + 0.25 * color.b;
    }

    float3 ToneMap(float3 color)
    {
        return color / (1 + Luminance(color));
    }

    float3 UnToneMap(float3 color)
    {
        return color / (1 - Luminance(color));
    }
    float2 GetClosestFragment(float2 uv)
    {
        float2 k = _CameraDepthTex_TexelSize.xy;
        //在上下左右四个点
        const float4 neighborhood = float4(
            SAMPLE_DEPTH_TEXTURE(_CameraDepthTex, sampler_PointClamp, clamp(uv - k,0,1)),
            SAMPLE_DEPTH_TEXTURE(_CameraDepthTex, sampler_PointClamp, clamp(uv + float2(k.x, -k.y),0,1)),
            SAMPLE_DEPTH_TEXTURE(_CameraDepthTex, sampler_PointClamp, clamp(uv + float2(-k.x, k.y),0,1)),
            SAMPLE_DEPTH_TEXTURE(_CameraDepthTex, sampler_PointClamp, clamp(uv + k,0,1))
        );
        // 获取离相机最近的点
        #if defined(UNITY_REVERSED_Z)
            #define COMPARE_DEPTH(a, b) step(b, a)
        #else
            #define COMPARE_DEPTH(a, b) step(a, b)
        #endif
        // 获取离相机最近的点，这里使用 lerp 是避免在shader中写分支判断
        float3 result = float3(0.0, 0.0, SAMPLE_DEPTH_TEXTURE(_CameraDepthTex, sampler_PointClamp, uv));
        result = lerp(result, float3(-1.0, -1.0, neighborhood.x), COMPARE_DEPTH(neighborhood.x, result.z));
        result = lerp(result, float3( 1.0, -1.0, neighborhood.y), COMPARE_DEPTH(neighborhood.y, result.z));
        result = lerp(result, float3(-1.0,  1.0, neighborhood.z), COMPARE_DEPTH(neighborhood.z, result.z));
        result = lerp(result, float3( 1.0,  1.0, neighborhood.w), COMPARE_DEPTH(neighborhood.w, result.z));
        return (uv + result.xy * k);
    }

  
    float2 historyPostion(float2 jitted_uv,float2 un_jitted_uv)
    {
        float depth = SAMPLE_TEXTURE2D_X(_CameraDepthTex, sampler_PointClamp, un_jitted_uv).r;
        #if UNITY_REVERSED_Z
                    depth = 1.0 - depth;
        #endif
                    depth = 2.0 * depth - 1.0;
        #if UNITY_UV_STARTS_AT_TOP
                    un_jitted_uv.y = 1.0f - un_jitted_uv.y;
        #endif
    
        float3 viewPos = ComputeViewSpacePosition(un_jitted_uv, depth, _I_P_Current_jittered);
        float4 worldPos = float4(mul(unity_CameraToWorld, float4(viewPos, 1.0)).xyz, 1.0);
        
        float4 historyNDC = mul(_PrevViewProjM_TAA, worldPos);
        historyNDC /= historyNDC.w;
        
        
        
        historyNDC.xy = historyNDC.xy * 0.5f + 0.5f;
       
        return historyNDC.xy;
    }
    float3 GethistoryPos(float2 un_jitted_uv)
    {
        float depth = SAMPLE_TEXTURE2D_X(_CameraDepthTex, sampler_PointClamp, un_jitted_uv).r;
        #if UNITY_REVERSED_Z
                    depth = 1.0 - depth;
        #endif
                    depth = 2.0 * depth - 1.0;
        #if UNITY_UV_STARTS_AT_TOP
                    un_jitted_uv.y = 1.0f - un_jitted_uv.y;
        #endif
    
        float3 viewPos = ComputeViewSpacePosition(un_jitted_uv+CurrentOffset, depth, _I_P_Current_jittered);
        float4 worldPos = float4(mul(unity_CameraToWorld, float4(viewPos, 1.0)).xyz, 1.0);
        float4 historyNDC = mul(_PrevViewProjM_TAA, worldPos);
        historyNDC /= historyNDC.w;
        
        //historyNDC.xy=historyNDC.xy+PreviousOffset;
        
        return historyNDC;
    }
    
    float3 RGBToYCoCg( float3 RGB )
    {
	    float Y  = dot( RGB, float3(  1, 2,  1 ) );
	    float Co = dot( RGB, float3(  2, 0, -2 ) );
	    float Cg = dot( RGB, float3( -1, 2, -1 ) );
	    
	    float3 YCoCg = float3( Y, Co, Cg );
	    return YCoCg;
    }

    float3 YCoCgToRGB( float3 YCoCg )
    {
	    float Y  = YCoCg.x * 0.25;
	    float Co = YCoCg.y * 0.25;
	    float Cg = YCoCg.z * 0.25;

	    float R = Y + Co - Cg;
	    float G = Y + Cg;
	    float B = Y - Co - Cg;

	    float3 RGB = float3( R, G, B );
	    return RGB;
    }
    float4 sample_colorAndToneMap(Texture2D<float4> tex, float2 uv)
    {
        float4 re=SAMPLE_TEXTURE2D_X(tex, sampler_LinearClamp, uv);
        re.rgb=RGBToYCoCg(ToneMap(re.rgb));
        return re;
        //
    }
    void minmax(in float2 uv, out float4 color_min, out float4 color_max, out float4 color_avg)
    {
        float2 du = float2(_SourceTex_TexelSize.x, 0.0);
        float2 dv = float2(0.0, _SourceTex_TexelSize.y);

        float4 ctl = sample_colorAndToneMap(_SourceTex, uv - dv - du);
        float4 ctc = sample_colorAndToneMap(_SourceTex, uv - dv);
        float4 ctr = sample_colorAndToneMap(_SourceTex, uv - dv + du);
        float4 cml = sample_colorAndToneMap(_SourceTex, uv - du);
        float4 cmc = sample_colorAndToneMap(_SourceTex, uv);
        float4 cmr = sample_colorAndToneMap(_SourceTex, uv + du);
        float4 cbl = sample_colorAndToneMap(_SourceTex, uv + dv - du);
        float4 cbc = sample_colorAndToneMap(_SourceTex, uv + dv);
        float4 cbr = sample_colorAndToneMap(_SourceTex, uv + dv + du);

        color_min = min(ctl, min(ctc, min(ctr, min(cml, min(cmc, min(cmr, min(cbl, min(cbc, cbr))))))));
        color_max = max(ctl, max(ctc, max(ctr, max(cml, max(cmc, max(cmr, max(cbl, max(cbc, cbr))))))));

        color_avg = (ctl + ctc + ctr + cml + cmc + cmr + cbl + cbc + cbr) / 9.0;
        
        // float2 ss_offset11 = float2(_SourceTex_TexelSize.x, _SourceTex_TexelSize.y);
        // float4 c00 =_SourceTex.Sample(sampler_SourceTex,uv-ss_offset11);
        // float4 c11 =_SourceTex.Sample(sampler_SourceTex,uv+ss_offset11);
        // color_min = min(c00, c11);
        // color_max = max(c00, c11);
        // color_avg = (c00 + c11) / 2.0;


    }
    float4 clip_aabb(float3 aabb_min, float3 aabb_max, float4 avg, float4 input_texel)
    {
        float3 p_clip = 0.5 * (aabb_max + aabb_min);
        float3 e_clip = 0.5 * (aabb_max - aabb_min) + FLT_EPS;
        float4 v_clip = input_texel - float4(p_clip, avg.w);
        float3 v_unit = v_clip.xyz / e_clip;
        float3 a_unit = abs(v_unit);
        float ma_unit = max(a_unit.x, max(a_unit.y, a_unit.z));

        if (ma_unit > 1.0)
            return float4(p_clip, avg.w) + v_clip / ma_unit;
        else
            return input_texel;
    }
    float3 ClipHistory(float3 History, float3 BoxMin, float3 BoxMax)
    {
        float3 Filtered = (BoxMin + BoxMax) * 0.5f;
        float3 RayOrigin = History;
        float3 RayDir = Filtered - History;
        RayDir = abs( RayDir ) < (1.0/65536.0) ? (1.0/65536.0) : RayDir;
        float3 InvRayDir = rcp( RayDir );
        float3 MinIntersect = (BoxMin - RayOrigin) * InvRayDir;
        float3 MaxIntersect = (BoxMax - RayOrigin) * InvRayDir;
        float3 EnterIntersect = min( MinIntersect, MaxIntersect );
        float ClipBlend = max( EnterIntersect.x, max(EnterIntersect.y, EnterIntersect.z ));
        ClipBlend = saturate(ClipBlend);
        return lerp(History, Filtered, ClipBlend);
    }

    ENDHLSL
    SubShader
    {
        // No culling or depth
        Cull Off ZWrite Off ZTest Always
        Pass
        {
            Name "TAA"
            HLSLPROGRAM
            #pragma vertex vert_TAA
            #pragma fragment frag_TAA
            v2f vert_TAA(appdata v)
            {
                v2f o;
                o.vertex=TransformObjectToHClip(v.vertex.xyz);
                o.uv=v.uv;
        
                return o;
            }
            half4 frag_TAA (v2f i) : SV_Target
            {
                //float2 closest=GetClosestFragment(i.uv);
                //float2 closest=(i.uv);
                float2 uv_jitted = UnityStereoTransformScreenSpaceTex(i.uv);
                float2 un_jitted_uv = uv_jitted - CurrentOffset;
                float4 current = _SourceTex.Sample(sampler_LinearClamp,un_jitted_uv);
                if(_IgnoreHistory)
                {
                    return current;
                }
                
                //float3 worldpos=GethistoryPos(un_jitted_uv);
                //return float4(worldpos,1.0f);
                float2 HistoryUV = historyPostion(uv_jitted,un_jitted_uv);

                if(HistoryUV.x < 0 || HistoryUV.y < 0 || HistoryUV.x > 1.0f || HistoryUV.y > 1.0f)
                {
                    return current;
                }
                float4 hisotry = _HistoryBuffer.Sample(sampler_LinearClamp,HistoryUV);
                hisotry.rgb=RGBToYCoCg(ToneMap(hisotry));
                current.rgb=RGBToYCoCg(ToneMap(current));
                
                
                
                float4 color_min, color_max, color_avg;
                minmax(un_jitted_uv, color_min, color_max, color_avg);
                
                hisotry=clip_aabb(color_min, color_max, color_avg, hisotry);
                //hisotry.rgb=ClipHistory(hisotry.rgb,color_min,color_max);
                hisotry.rgb = UnToneMap(YCoCgToRGB(hisotry));
                current.rgb = UnToneMap(YCoCgToRGB(current));
                return lerp(hisotry, current, _BlendAlpha);
            }
            ENDHLSL
        }
        Pass
        {
            Name "TAA Tonemapping"
            HLSLPROGRAM
            #pragma vertex vert_ToneMapping
            #pragma fragment frag_Tonemapping
            v2f vert_ToneMapping(appdata v)
            {
                v2f o;
                o.vertex=TransformObjectToHClip(v.vertex.xyz);
                o.uv=v.uv;
        
                return o;
            }
            half4 frag_Tonemapping (v2f i) : SV_Target
            {
                float4 col=_SourceTex.Sample(sampler_LinearClamp,i.uv);
                col.rgb=ToneMap(col.rgb);
                return col;
            }
            ENDHLSL
        }
    }
}